# -*- coding: utf-8 -*-
'''
Created on May 20, 2012

@author: tel
'''
import sys, os, time, numpy, site, re
sys.path.append('/home/tel/workspace/pmcce/src')
#site.addsitedir('/home/tel/workspace/pmcce/src')
from collections import namedtuple

import makedb.levels as levels
from histo import SixPanelHisto
from scatter import Scatter
from hexbin import Hexbin
from bezier import Bezier
from sqlalchemy import Table, Column, Integer, ForeignKey, create_engine, event, MetaData
from sqlalchemy.orm import relationship, backref, aliased, sessionmaker, joinedload
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.pool import NullPool

def _fk_pragma_on_connect(dbapi_con, con_record):
    dbapi_con.execute('pragma foreign_keys=ON')

def LoadSession(db):
    Session = sessionmaker(bind=db)
    session = Session()
    return session
    
def Nearby_Mutants(session):
    for prot, res in session.query(levels.ProtData, levels.ResData).filter(levels.ResData.seqadv!='').yield_per(10):
        near_res = {}
        for atom in res.Atoms:
            if atom.atom=='CB':
                for near_atom in prot.Atoms:
                    if atom.Dist(near_atom) < 10 and near_atom.atom=='CB':
                        near_res[str(near_atom.Res.res)+str(near_atom.Res.chainid)+str(near_atom.Res.resn)] = near_atom.Res
        print 'this mutant:'
        print '%s/e%s\t%s\t%s\t%s\t' % (prot.pdb,res.Run.EPSILON_PROT,res.res,res.chainid,res.resn)
        print 'is within 10 A of these other residues:'
        for r in near_res.values():
            print '%s/e%s\t%s\t%s\t%s\t' % (prot.pdb,res.Run.EPSILON_PROT,r.res,r.chainid,r.resn)

#for prot,run,res in session.query(levels.ProtData, levels.RunData, levels.ResData).join(levels.ProtData.Runs).join(levels.RunData.Ress).filter(levels.ResData.charge!=None).yield_per(10): 
def SumCrgAt(session, titr):
    '''
    gives the ionization of every residue at ph/em = titr
    '''
    print levels.ResData.StrHeaders(), '\tSum_Crg_at_%.3f' % titr
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        print res, '\t%.3f' % res.SumCrgAt(titr)
        
def DPKGreaterThan(session, dpk):
    '''
    gives every residue for which the pKa differs from the solution value by dpk or more, and the pKa
    '''
    print levels.ResData.StrHeaders(), '\tpKa\tDpKa'
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        if abs(res.dpk)>dpk:
            print res, '\t%.3f\t%.3f' % (res.__getattribute__('pKa/Em'),res.dpk)
            
def DsolGreaterThan(session, dsol):
    '''
    gives every residue for which the pKa differs from the solution value by dpk or more, and the pKa
    '''
    print levels.ResData.StrHeaders(), '\tDsol'
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        if abs(res.dsol)>dsol:
            print res, '\t%.3f' % res.dsol

def RMSDGreaterThan(session, rmsd):
    print levels.ResData.StrHeaders(), '\tRMSD'
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        titrs = sorted(res.Confs[1].occupancy.keys())
        initial = res.MostOccAt(titrs[0])
        final = res.MostOccAt(titrs[-1])
        RMSD = initial.RMSD(final,session)
        if RMSD>rmsd:
            print res, '\t%.3f' % RMSD
            
def CrystalConformers(session):
    print levels.ResData.StrHeaders(), '\tConformers_with_Crystal_Heavy_Atoms'
    for res in session.query(levels.ResData).yield_per(10):
        crysconfs = [1]
        for i,conf in enumerate(res.Confs[2:]):
            if res.Confs[1].RMSD(conf,session)==0:
                crysconfs.append(i+2)
        print res, '\t%s' % crysconfs

def SumCrgNonmonotonic(session, tolerance=.05):
    '''
    gives every residue for which the ionization goes up and down across the titation, and the ionization at all titration points
    will not report as nonmonotonic if variation is within the tolerance
    '''
    print levels.ResData.StrHeaders()
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        sumcrg = sorted(res.sumcrg.items()) #sumcrg is stored as a dict
        ionization = zip(*sumcrg)[1]
        deltas = map(lambda x,y: y-x, ionization[:-1], ionization[1:])
        big_pos = False
        big_neg = False
        for delta in deltas:
            if delta>tolerance:
                if big_neg==True:
                    print res
                    print ionization
                    break
                else:
                    big_pos=True
            elif delta<-tolerance:
                if big_pos==True:
                    print res
                    print ionization
                    break
                else:
                    big_neg=True
                    
def AtomInfo(session, res, resn, chainid, atomn):
    atom = session.query(levels.AtomData).join(levels.ResData).filter(levels.ResData.res==res).filter(levels.ResData.resn==resn).filter(levels.ResData.chainid==chainid).filter(levels.AtomData.atomn==atomn).first()
    print atom.StrHeaders(), 'X\tY\tZ'
    print atom, ('%.3f\t%.3f\t%.3f\t' % (atom.x, atom.y, atom.z))
            
if __name__=='__main__':
    funcdict = {'SumCrgAt':SumCrgAt,'DPKGreaterThan':DPKGreaterThan,'DsolGreaterThan':DsolGreaterThan,'RMSDGreaterThan':RMSDGreaterThan,'CrystalConformers':CrystalConformers,'SumCrgNonmonotonic':SumCrgNonmonotonic,'AtomInfo':AtomInfo}
    try:
        dbpath = sys.argv[1]
        func = funcdict[sys.argv[2]]
    except (KeyError, IndexError):
        print 'Usage:'
        print 'path-to-database function [argument]\n'
        print 'Possible functions:'
        for key in funcdict.keys():
            print key
        sys.exit(0)
    arg = sys.argv[3:4]
    dir = '/home/tel/test/db1'
    db = create_engine('sqlite:///' + dbpath, poolclass=NullPool)
    event.listen(db, 'connect', _fk_pragma_on_connect)
    meta= MetaData(db)
    
    for level in levels.hierarchy:
        level.Table(meta=meta, db=db, reloadt=True)
    session = LoadSession(db)
    
    #Nearby_Mutants(session)
    if arg:
        func(session, float(arg[0]))
    else:
        func(session)
